package com.etl.component.postgresql;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.google.common.collect.Lists;
import com.component.api.constants.ComponentCodeConstants;
import com.component.api.constants.ComponentConstants;
import com.component.api.excepiton.ComponentException;
import com.component.api.model.PortData;
import com.component.api.model.ProcessResult;
import com.component.api.model.data.SetData;
import com.component.api.model.dto.DatabaseParamDTO;
import com.component.api.model.param.ParamPair;
import com.component.api.model.param.TableColumn;
import com.component.api.utils.SqlUtils;
import com.component.api.utils.TypesUtils;
import com.common.utils.constant.SQLConstants;
import com.etl.component.common.AbstractFlinkComponent;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.DataSet;
import org.apache.flink.api.java.io.jdbc.JDBCOutputFormat;
import org.apache.flink.api.java.typeutils.RowTypeInfo;
import org.apache.flink.types.Row;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * 描述：
 * postgresql数据库读取
 *  因为是单表 都采用jdbc读取
 * @author xianggj
 * @Date 2021/10/12 18:10
 **/
public class PostgresqlWriteComponent extends AbstractFlinkComponent {

    /**
     * 主键
     */
    private String colKey;

    @Override
    public ProcessResult process(List<PortData> datas, List<ParamPair> paramPairs) {
        Map<String, Object> params = paramToObject(paramPairs);
        DatabaseParamDTO databaseParamDTO = (DatabaseParamDTO)params.get(ComponentConstants.CONFIG);
        String driver = databaseParamDTO.getDriver();
        String userName = databaseParamDTO.getUserName();
        String pwd = databaseParamDTO.getPwd();
//        String schema = databaseParamDTO.getSchema() == null? "public": databaseParamDTO.getSchema();
        String url = databaseParamDTO.getUrl() ;
        String tableName = databaseParamDTO.getTableName();
        String writeModel = (String)params.get(ComponentConstants.WRITE_MODEL);

        SetData input = (SetData)datas.get(0).getValue();
        //列从输入dataset里获取有问题 最好从参数获取

        List<TableColumn> columns = input.getColumns();
        for (TableColumn tc : columns) {
            if (tc.getColumnKey()){
                this.colKey = tc.getColumnName();
                break;
            }
        }
        if (this.colKey == null){
            this.colKey = SQLConstants.ID;
        }
        DataSet<Row> result = (DataSet<Row>) input.getDataset();

        int[] fieldTypes =new int[columns.size()];
//        String sql = getSqlByModel(tableName, writeModel, columns, fieldTypes);
        RowTypeInfo type = (RowTypeInfo) result.getType();
        TypeInformation[] fieldTypeInfos = type.getFieldTypes();
        String[] fieldNames = type.getFieldNames();
        //插入更新 进行数据分流
        if (ComponentConstants.INSERT_UPDATE.equalsIgnoreCase(writeModel)){
            List<String> ids = (List<String>)params.get(ComponentConstants.IDS);
            //数据分流
            if (CollUtil.isEmpty(ids)){
                result = execWrite(driver, url, userName,
                        pwd,tableName, ComponentConstants.APPEND, input,
                        result, fieldTypes,
                        fieldTypeInfos,
                        fieldNames);
                return execute(result, columns, "执行pg插入操作");
            } else {
                List<String> colName = input.getColName();
                final int index = indexOfIgnoreCase(colName, this.colKey);
                DataSet<Row> result2 = result.filter(e->{
                    Object field = e.getField(index);
                    if (field == null){
                        return false;
                    }
                    return !ids.contains(field.toString());
                });
                DataSet<Row> dataSet2 = execWrite(driver, url, userName,
                        pwd, tableName, ComponentConstants.APPEND, input,
                        result2, fieldTypes,
                        fieldTypeInfos,
                        fieldNames);
                try {
                    dataSet2.collect();
                } catch (Exception e) {
                    e.printStackTrace();
                    throw new ComponentException("执行插入时出现错误：", e);
                }
                DataSet<Row> result1 = result.filter(e->{
                    Object field = e.getField(index);
                    if (field == null){
                        return false;
                    }
                    return ids.contains(field.toString());
                });
                DataSet<Row> dataSet1 = execWrite(driver, url, userName,
                        pwd,tableName, ComponentConstants.UPDATE, input,
                        result1, fieldTypes,
                        fieldTypeInfos,
                        fieldNames);
                try {
                    dataSet1.collect();
                } catch (Exception e) {
                    e.printStackTrace();
                    throw new ComponentException("执行更新时出现错误：", e);
                }


                return execute(result, columns, "执行pg写入操作");
            }
        } else {
            result = execWrite(driver, url, userName,
                    pwd,tableName, writeModel, input,
                    result, fieldTypes,
                    fieldTypeInfos,
                    fieldNames);

            return execute(result, columns, "执行pg写入操作");
        }
    }

    private DataSet<Row>  execWrite(String driver, String url,
                                    String userName, String pwd, String tableName,
                                    String writeModel, SetData input,
                                    DataSet<Row> result,
                                    int[] fieldTypes,
                                    TypeInformation[] fieldTypeInfos,
                                    String[] fieldNames) {
        List<TableColumn> columns = input.getColumns();
        String sql = getSqlByModel(tableName, writeModel, fieldTypeInfos,
                fieldNames, fieldTypes);
        //update 特殊处理 将数据位置换一下
        if (ComponentConstants.UPDATE.equalsIgnoreCase(writeModel)){
            List<String> colName = input.getColName();
            int index = indexOfIgnoreCase(colName, this.colKey);
            result = result.map(e->{
                int arity = e.getArity();
                if (index == arity -1){
                    return e;
                }
                Object field = e.getField(index);
                Object endfield = e.getField(arity - 1);
                e.setField(arity - 1, field );
                e.setField(index, endfield );
                return e;
            });
        }

        //写数据
        JDBCOutputFormat jdbcOutput = JDBCOutputFormat.buildJDBCOutputFormat()
                .setDrivername(driver)
                .setDBUrl(url)
                .setUsername(userName)
                .setPassword(pwd)
                .setQuery(sql)
                .setSqlTypes(fieldTypes)
                .setBatchInterval(10)//设置10条数据提交一次
                .finish();
        result.output(jdbcOutput);
        return result;
    }


    private String getSqlByModel(String tableName, String writeModel, List<TableColumn>
            columns, int[] fieldTypes) {
        String sql = null;
        List<String> columnNames = new ArrayList<>();
        //添加策略 id重复问题需要自行解决
        if (ComponentConstants.APPEND.equalsIgnoreCase(writeModel)){
            int i = 0;
            for (TableColumn column : columns) {
                fieldTypes[i] = TypesUtils.typeStrToInt(column.getDataType());
                columnNames.add(column.getColumnName());
                i++;
            }
            sql = createInsertSQL(tableName, columnNames);
            //修改策略 默认根据id修改 如果id不存在跳过
        } else if (ComponentConstants.UPDATE.equalsIgnoreCase(writeModel)){
            TableColumn endColumns = columns.get(columns.size() - 1);
            //id为最后一列就不循环
            if (this.colKey.equalsIgnoreCase(endColumns.getColumnName())){
                fieldTypes[columns.size()-1] = TypesUtils.typeStrToInt(endColumns.getDataType());
            }
            for (int j = 0; j < columns.size() -1 ; j++) {
                TableColumn column = columns.get(j);
                if (this.colKey.equalsIgnoreCase(column.getColumnName())){
                    fieldTypes[columns.size()-1] = TypesUtils.typeStrToInt(column.getDataType());
                    column = endColumns;
                }
                fieldTypes[j] = TypesUtils.typeStrToInt(column.getDataType());
                columnNames.add(column.getColumnName());
            }
            sql = createUpdateSQL(tableName, columnNames);
            //覆盖
        } else {
            int i = 0;
            for (TableColumn column : columns) {
                fieldTypes[i] = TypesUtils.typeStrToInt(column.getDataType());
                columnNames.add(column.getColumnName());
                i++;
            }
            sql = createReplaceSQL(tableName, columnNames);
        }
        return sql;
    }


    private String getSqlByModel(String tableName, String writeModel,
                                 TypeInformation[] fieldTypeInfos, String[]
            columnNames, int[] fieldTypes) {
        String sql = null;
        //添加策略 id重复问题需要自行解决
        if (ComponentConstants.APPEND.equalsIgnoreCase(writeModel)){
            int i = 0;
            for (TypeInformation TypeInfo : fieldTypeInfos) {
                fieldTypes[i] = typeToInt(TypeInfo);
                i++;
            }
            sql = createInsertSQL(tableName, Arrays.asList(columnNames));
            //修改策略 默认根据id修改 如果id不存在跳过
        } else if (ComponentConstants.UPDATE.equalsIgnoreCase(writeModel)){
            String endColumns = columnNames[columnNames.length - 1];
            TypeInformation endTypes =fieldTypeInfos[fieldTypeInfos.length - 1];
            //id为最后一列就不循环
            if (this.colKey.equalsIgnoreCase(endColumns)){
                fieldTypes[columnNames.length-1] = TypesUtils.typeStrToInt(endColumns);
            }
            List<String> colNames = new ArrayList<>();
            for (int j = 0; j < columnNames.length -1 ; j++) {
                TypeInformation fieldTypeInfo = fieldTypeInfos[j];
                if (this.colKey.equalsIgnoreCase(columnNames[j])){
                    //交换位置
                    fieldTypes[fieldTypeInfos.length - 1] = typeToInt(fieldTypeInfo);
                    columnNames[j] = endColumns;
                    fieldTypeInfo = endTypes;
                }
                colNames.add(columnNames[j]);
                fieldTypes[j] = typeToInt(fieldTypeInfo);
            }
            sql = createUpdateSQL(tableName, colNames);
            //覆盖
        } else {
            int i = 0;
            for (TypeInformation TypeInfo : fieldTypeInfos) {
                fieldTypes[i] = typeToInt( TypeInfo);
                i++;
            }
            sql = createReplaceSQL(tableName, Arrays.asList(columnNames));
        }
        return sql;
    }

    /**
     * 生成创建语句
     * @param tableName
     * @param columnNames
     * @return
     */
    private String createInsertSQL(String tableName, List<String> columnNames) {
        StringBuffer insertSql = new StringBuffer(SQLConstants.INSERT_INFO);
        insertSql.append(SqlUtils.addEscape(tableName , SQLConstants.PG_ESC));
        insertSql.append(SQLConstants.LEFT_BRACKET);
        for (int i = 0; i < columnNames.size() ; ++i) {
            if (i > 0){
                insertSql.append(SQLConstants.COMMA);
            }
            insertSql.append(SqlUtils.addEscape(columnNames.get(i) , SQLConstants.PG_ESC));
        }
        insertSql.append(SQLConstants.RIGHT_BRACKET);
        insertSql.append(SQLConstants.VALUES);
        insertSql.append(SQLConstants.LEFT_BRACKET);
        String[] array = new String[columnNames.size()];
        Arrays.fill(array, SQLConstants.QUESTION);
        insertSql.append(StrUtil.join(SQLConstants.COMMA, array));
        insertSql.append(SQLConstants.RIGHT_BRACKET);
        return insertSql.toString();
    }

    /**
     * 生成覆盖语句
     * @param tableName
     * @param columnNames
     * @return
     */
    private String createReplaceSQL(String tableName, List<String> columnNames) {
        //insert into test values (1,'test',now()) on conflict (id) do update set info=excluded.info,crt_time=excluded.crt_time;
        StringBuffer insertSql = new StringBuffer(SQLConstants.INSERT_INFO);
        insertSql.append(SqlUtils.addEscape(tableName , SQLConstants.PG_ESC));
        insertSql.append(SQLConstants.LEFT_BRACKET);
        for (int i = 0; i < columnNames.size() ; ++i) {
            if (i > 0){
                insertSql.append(SQLConstants.COMMA);
            }
            insertSql.append(SqlUtils.addEscape(columnNames.get(i) , SQLConstants.PG_ESC));
        }
        insertSql.append(SQLConstants.RIGHT_BRACKET);
        insertSql.append(SQLConstants.VALUES);
        insertSql.append(SQLConstants.LEFT_BRACKET);
        String[] array = new String[columnNames.size()];
        Arrays.fill(array, SQLConstants.QUESTION);
        insertSql.append(StrUtil.join(SQLConstants.COMMA, array));
        insertSql.append(SQLConstants.RIGHT_BRACKET);
        //拼装conflict
//        insertSql.append(SQLConstants.ON);
//        insertSql.append(SQLConstants.CONFLICT);
//        insertSql.append(SQLConstants.LEFT_BRACKET);
//        insertSql.append(SqlUtils.addEscape(this.colKey, SQLConstants.PG_ESC));
//        insertSql.append(SQLConstants.RIGHT_BRACKET);
//        insertSql.append(SQLConstants.DO);
//        insertSql.append(SQLConstants.UPDATE);
//        insertSql.append(SQLConstants.SET);
//        for (String col: columnNames) {
//            insertSql.append(col);
//            insertSql.append(SQLConstants.EQ);
//            insertSql.append(SQLConstants.QUESTION);
//            insertSql.append(SQLConstants.COMMA);
//        }
//        insertSql.deleteCharAt(insertSql.length()-1);
        return insertSql.toString();
    }
    /**
     * 生成修改语句
     * @param tableName
     * @param columnNames
     * @return
     */
    private String createUpdateSQL(String tableName, List<String> columnNames) {
        //UPDATE `test`.`test` SET `name` = 'a', `age` = 1 WHERE `id` = 1;
        StringBuffer insertSql = new StringBuffer(SQLConstants.UPDATE);
        insertSql.append(SqlUtils.addEscape(tableName , SQLConstants.PG_ESC));
        insertSql.append(SQLConstants.SET);
        for (String col: columnNames) {
            insertSql.append(SqlUtils.addEscape(col , SQLConstants.PG_ESC));
            insertSql.append(SQLConstants.EQ);
            insertSql.append(SQLConstants.QUESTION);
            insertSql.append(SQLConstants.COMMA);
        }
        insertSql.deleteCharAt(insertSql.length()-1);
        insertSql.append(SQLConstants.WHERE);
        insertSql.append(this.colKey);
        insertSql.append(SQLConstants.EQ);
        insertSql.append(SQLConstants.QUESTION);

        return insertSql.toString();
    }


    @Override
    public List<PortData> previewOutputColumns(List<PortData> datas, List<ParamPair> paramPairs) {
        return Lists.newArrayList();
    }

    @Override
    public String getRule() {
        return ComponentCodeConstants.PG_OUTPUT;
    }

}
